import cv2
import numpy as np
from PIL import Image
from rembg import remove
import mediapipe as mp
import time
from enum import Enum

class BackgroundRemovalMethod(Enum):
    REMBG = "rembg"
    MEDIAPIPE = "mediapipe"
    DEEPLAB = "deeplab"

class IDPhotoMaker:
    def __init__(self, method=BackgroundRemovalMethod.REMBG):
        """初始化，支持多种背景去除方法"""
        self.method = method
        
        # 预设的证件照尺寸
        self.PHOTO_SIZES = {
            '1inch': (295, 413),      # 一寸
            '2inch': (413, 579),      # 二寸
            'small1inch': (260, 378),  # 小一寸
            'small2inch': (413, 531),  # 小二寸
        }
        # 默认使用一寸照片尺寸
        self.ID_PHOTO_WIDTH = self.PHOTO_SIZES['1inch'][0]
        self.ID_PHOTO_HEIGHT = self.PHOTO_SIZES['1inch'][1]
        
        # 根据选择的方法初始化相应模型
        if method == BackgroundRemovalMethod.MEDIAPIPE:
            self.mp_selfie_segmentation = mp.solutions.selfie_segmentation
            self.selfie_seg = self.mp_selfie_segmentation.SelfieSegmentation(model_selection=1)
        elif method == BackgroundRemovalMethod.DEEPLAB:
            try:
                import torch
                import torchvision.transforms as transforms
                self.model = torch.hub.load('pytorch/vision:v0.10.0', 
                                          'deeplabv3_resnet50', 
                                          pretrained=True)
                self.model.eval()
                if torch.cuda.is_available():
                    self.model = self.model.cuda()
                self.transform = transforms.Compose([
                    transforms.ToTensor(),
                    transforms.Normalize(
                        mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225]
                    ),
                ])
            except ImportError:
                raise ImportError(
                    "使用 DeepLab 方法需要安装 PyTorch。\n"
                    "请运行: pip install torch torchvision"
                )

    def remove_background(self, image_path):
        """根据选择的方法调用相应的背景去除函数"""
        if self.method == BackgroundRemovalMethod.REMBG:
            return self._remove_background_rembg(image_path)
        elif self.method == BackgroundRemovalMethod.MEDIAPIPE:
            return self._remove_background_mediapipe(image_path)
        elif self.method == BackgroundRemovalMethod.DEEPLAB:
            return self._remove_background_deeplab(image_path)

    def _remove_background_rembg(self, image_path):
        """使用 rembg 去除背景"""
        input_image = Image.open(image_path)
        output_image = remove(
            input_image,
            alpha_matting=True,
            alpha_matting_foreground_threshold=250,
            alpha_matting_background_threshold=3,
            alpha_matting_erode_size=20
        )
        
        opencv_image = cv2.cvtColor(np.array(output_image), cv2.COLOR_RGBA2BGRA)
        b, g, r, a = cv2.split(opencv_image)
        image = cv2.merge([b, g, r])
        mask = a / 255.0
        
        # 后处理
        kernel = np.ones((3,3), np.uint8)
        mask = cv2.erode(mask, kernel, iterations=1)
        mask = cv2.dilate(mask, kernel, iterations=1)
        _, mask = cv2.threshold(mask, 0.5, 1, cv2.THRESH_BINARY)
        
        return image, mask

    def _remove_background_mediapipe(self, image_path):
        """使用 MediaPipe 去除背景，增强肩部处理"""
        # 读取图片
        image = cv2.imread(image_path)
        
        # 预处理：调整图片大小
        max_size = 1500
        h, w = image.shape[:2]
        scale = min(max_size / max(h, w), 1.0)
        
        if scale < 1.0:
            new_size = (int(w * scale), int(h * scale))
            image = cv2.resize(image, new_size, interpolation=cv2.INTER_LANCZOS4)
            print(f"图片已优化: {(w,h)} -> {new_size}")
        
        # 转换为 RGB 格式
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 使用 MediaPipe 进行分割
        results = self.selfie_seg.process(image_rgb)
        mask = results.segmentation_mask
        
        # 增强 mask 处理
        # 1. 创建不同大小的 kernel
        kernel_small = np.ones((2,2), np.uint8)
        kernel_medium = np.ones((3,3), np.uint8)
        kernel_large = np.ones((7,7), np.uint8)  # 增大kernel用于更平滑的肩部处理
        
        # 创建椭圆形核，更适合处理肩部曲线
        ellip_kernel_small = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5))
        ellip_kernel_large = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9,9))
        
        # 2. 对原始 mask 进行平滑处理
        mask_original = cv2.GaussianBlur(mask, (7, 7), 0)  # 增加初始平滑
        
        # 3. 创建多层 mask
        # 主体 mask（较严格的阈值）
        _, mask_main = cv2.threshold(mask_original, 0.35, 1, cv2.THRESH_BINARY)
        # 肩部 mask（专门处理肩部）
        _, mask_shoulder = cv2.threshold(mask_original, 0.25, 1, cv2.THRESH_BINARY)
        # 边缘 mask（处理过渡区域）
        _, mask_edge = cv2.threshold(mask_original, 0.15, 1, cv2.THRESH_BINARY)
        # 细节 mask（保留细节）
        _, mask_detail = cv2.threshold(mask_original, 0.1, 1, cv2.THRESH_BINARY)
        
        # 4. 分别处理各层
        # 处理主体
        mask_main = cv2.morphologyEx(mask_main, cv2.MORPH_CLOSE, kernel_medium)
        
        # 处理肩部（重点优化）
        # 使用椭圆核进行闭运算
        mask_shoulder = cv2.morphologyEx(mask_shoulder, cv2.MORPH_CLOSE, ellip_kernel_large)
        # 使用大核心进行平滑
        mask_shoulder = cv2.morphologyEx(mask_shoulder, cv2.MORPH_CLOSE, kernel_large)
        # 再次使用椭圆核处理
        mask_shoulder = cv2.morphologyEx(mask_shoulder, cv2.MORPH_CLOSE, ellip_kernel_small)
        
        # 处理边缘过渡区域
        mask_edge = cv2.morphologyEx(mask_edge, cv2.MORPH_CLOSE, ellip_kernel_small)
        mask_edge = cv2.GaussianBlur(mask_edge, (5, 5), 0)
        
        # 5. 边缘检测和平滑
        # 使用 Sobel 算子检测边缘
        sobelx = cv2.Sobel(mask_main, cv2.CV_64F, 1, 0, ksize=5)  # 增大 ksize
        sobely = cv2.Sobel(mask_main, cv2.CV_64F, 0, 1, ksize=5)
        edges = np.sqrt(sobelx**2 + sobely**2)
        edges = cv2.normalize(edges, None, 0, 1, cv2.NORM_MINMAX)
        
        # 平滑边缘区域
        edges = cv2.dilate(edges, kernel_large, iterations=2)
        edges = cv2.GaussianBlur(edges, (7, 7), 0)
        
        # 6. 智能合并
        # 首先合并主体和肩部
        mask_combined = np.where(edges > 0.1, mask_shoulder, mask_main)
        
        # 添加边缘过渡
        mask_with_edge = cv2.addWeighted(mask_combined, 0.7, mask_edge, 0.3, 0)
        
        # 添加细节
        mask_final = cv2.addWeighted(mask_with_edge, 0.8, mask_detail, 0.2, 0)
        
        # 7. 最终平滑处理
        # 使用较大的高斯核平滑
        mask_final = cv2.GaussianBlur(mask_final, (5, 5), 0)
        # 使用椭圆核进行最后的形态学处理
        mask_final = cv2.morphologyEx(mask_final, cv2.MORPH_CLOSE, ellip_kernel_small)
        # 最终的边缘平滑
        mask_final = cv2.GaussianBlur(mask_final, (3, 3), 0)
        
        return image, mask_final.astype(np.float32)

    def _remove_background_deeplab(self, image_path):
        """使用 DeepLab 去除背景，增强边缘处理"""
        try:
            import torch
            import torchvision.transforms as transforms
        except ImportError:
            raise ImportError(
                "使用 DeepLab 方法需要安装 PyTorch。\n"
                "请运行: pip install torch torchvision"
            )
            
        # 读取图片
        image = cv2.imread(image_path)
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        input_image = Image.fromarray(image_rgb)
        
        # 预处理
        input_tensor = self.transform(input_image).unsqueeze(0)
        if torch.cuda.is_available():
            input_tensor = input_tensor.cuda()
        
        # 预测
        with torch.no_grad():
            output = self.model(input_tensor)['out'][0]
        
        # 获取人像mask (类别15是人)
        mask = (output.argmax(0) == 15).cpu().numpy().astype(np.float32)
        
        # 增强mask处理
        # 1. 创建不同大小的kernel
        kernel_small = np.ones((2,2), np.uint8)
        kernel_medium = np.ones((3,3), np.uint8)
        kernel_large = np.ones((5,5), np.uint8)
        
        # 2. 对原始mask进行平滑处理
        mask_original = cv2.GaussianBlur(mask, (3, 3), 0)
        
        # 3. 创建三层mask
        # 主体mask（严格阈值）
        _, mask_main = cv2.threshold(mask_original, 0.8, 1, cv2.THRESH_BINARY)
        # 头发和边缘mask（中等阈值）
        _, mask_hair = cv2.threshold(mask_original, 0.3, 1, cv2.THRESH_BINARY)
        # 细节mask（宽松阈值）
        _, mask_detail = cv2.threshold(mask_original, 0.1, 1, cv2.THRESH_BINARY)
        
        # 4. 分别处理各层
        # 处理主体
        mask_main = cv2.erode(mask_main, kernel_medium, iterations=1)
        mask_main = cv2.dilate(mask_main, kernel_medium, iterations=1)
        
        # 处理头发层
        mask_hair = cv2.morphologyEx(mask_hair, cv2.MORPH_CLOSE, kernel_small)  # 闭运算填充小孔
        mask_hair = cv2.dilate(mask_hair, kernel_small, iterations=2)
        
        # 处理细节层
        mask_detail = cv2.dilate(mask_detail, kernel_small, iterations=1)
        
        # 5. 边缘检测和增强
        edges = cv2.Canny(mask_main.astype(np.uint8) * 255, 100, 200)
        edges = cv2.dilate(edges, kernel_large, iterations=2)  # 扩大边缘区域
        edges = edges.astype(np.float32) / 255.0
        
        # 6. 智能合并
        # 首先合并主体和头发
        mask_combined = np.where(edges > 0, mask_hair, mask_main)
        
        # 在边缘区域添加细节
        edge_detail = edges * mask_detail
        mask_final = cv2.addWeighted(mask_combined, 0.7, edge_detail, 0.3, 0)
        
        # 7. 最终处理
        # 对比度增强
        mask_final = cv2.normalize(mask_final, None, 0, 1, cv2.NORM_MINMAX)
        
        # 平滑边缘
        mask_final = cv2.GaussianBlur(mask_final, (3, 3), 0)
        
        # 再次增强对比度
        _, mask_final = cv2.threshold(mask_final, 0.2, 1, cv2.THRESH_BINARY)
        mask_final = cv2.GaussianBlur(mask_final, (3, 3), 0)
        
        # 8. 最终的边缘羽化
        mask_final = cv2.erode(mask_final, kernel_small, iterations=1)
        mask_final = cv2.GaussianBlur(mask_final, (5, 5), 0)
        
        return image, mask_final

    def change_background(self, image, mask, color):
        """更改背景颜色，增强边缘处理"""
        h, w = image.shape[:2]
        
        # 创建背景
        if color == 'gradient':
            background = np.zeros((h, w, 3), np.uint8)
            for i in range(h):
                ratio = i / h
                background[i, :] = [
                    int(219 * (1 - ratio) + 255 * ratio),
                    int(142 * (1 - ratio) + 255 * ratio),
                    int(67 * (1 - ratio) + 255 * ratio)
                ]
        else:
            # 创建纯色背景
            color_dict = {
                'red': [0, 0, 255],
                'blue': [219, 142, 67],
                'white': [255, 255, 255]
            }
            background = np.full((h, w, 3), color_dict[color], dtype=np.uint8)
        
        # 增强边缘羽化
        blur_radius = 1  # 减小羽化半径以保持头发细节
        mask_blurred = cv2.GaussianBlur(mask, (blur_radius*2+1, blur_radius*2+1), blur_radius)
        
        # 合并前景和背景
        mask_3channel = np.stack([mask_blurred] * 3, axis=-1)
        foreground = (image.astype(np.float32) * mask_3channel)
        bg = (background.astype(np.float32) * (1 - mask_3channel))
        result = foreground + bg
        
        return result.astype(np.uint8)

    def resize_to_id_photo(self, image):
        """调整图片尺寸为标准一寸照大小，优化裁剪方式"""
        h, w = image.shape[:2]
        target_ratio = self.ID_PHOTO_WIDTH / self.ID_PHOTO_HEIGHT
        current_ratio = w / h

        if current_ratio > target_ratio:
            # 图片太宽，需要裁剪宽度
            new_w = int(h * target_ratio)
            start_x = (w - new_w) // 2
            image = image[:, start_x:start_x+new_w]
        else:
            # 图片太高，需要裁剪高度
            new_h = int(w / target_ratio)
            start_y = (h - new_h) // 2
            image = image[start_y:start_y+new_h, :]

        # 调整到标准尺寸
        final_image = cv2.resize(image, (self.ID_PHOTO_WIDTH, self.ID_PHOTO_HEIGHT))
        return final_image

    def process_photo(self, input_path, output_path, bg_color='white'):
        """处理照片的主函数，添加详细的时间统计"""
        timing = {}  # 用于存储各阶段的时间
        
        # 1. 背景去除
        start = time.time()
        image, mask = self.remove_background(input_path)
        timing['背景去除'] = time.time() - start
        
        # 2. 更换背景颜色
        start = time.time()
        image_with_new_bg = self.change_background(image, mask, bg_color)
        timing['背景更换'] = time.time() - start
        
        # 3. 调整尺寸
        start = time.time()
        final_image = self.resize_to_id_photo(image_with_new_bg)
        timing['尺寸调整'] = time.time() - start
        
        # 4. 保存结果
        start = time.time()
        cv2.imwrite(output_path, final_image)


        timing['保存图片'] = time.time() - start
        
        # 返回时间统计
        return timing 